package com.at.sql16;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;

import static org.apache.flink.table.api.Expressions.$;

/**
 * 自定义 UDF 聚合函数
 * TODO 需求： 从学生的分数表ScoreTable中计算每个学生的加权平均分
 *
 * @author huangchao E-mail:fengquan8866@163.com
 * @version 创建时间：2024/10/6 21:58
 */
public class MyAggregateFunctionDemo5 {
    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStreamSource<Tuple3<String, Integer, Integer>> scoreWeightDS = env.fromElements(
                Tuple3.of("zs", 80, 3),
                Tuple3.of("zs", 90, 4),
                Tuple3.of("zs", 95, 5),
                Tuple3.of("ls", 81, 3),
                Tuple3.of("ls", 91, 4),
                Tuple3.of("ls", 96, 5)
        );

        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        Table scoreWeightTable = tableEnv.fromDataStream(scoreWeightDS,
                $("f0").as("name"),
                $("f1").as("score"),
                $("f2").as("weight"));
        tableEnv.createTemporaryView("scores", scoreWeightTable);

        // TODO 2.注册函数
        tableEnv.createTemporaryFunction("WeightedAvg", WeightedAvg.class);

        // TODO 3.调用 自定义函数
        tableEnv
                .sqlQuery("select name,WeightedAvg(score,weight) from scores group by name")
                .execute()
                .print();
    }

    // TODO 1.继承 AggregateFunction<返回类型, 累加器类型<加权总和，权重总和>>
    public static class WeightedAvg extends AggregateFunction<Double, Tuple2<Integer, Integer>> {

        @Override
        public Double getValue(Tuple2<Integer, Integer> integerIntegerTuple2) {
            return integerIntegerTuple2.f0 * 1D / integerIntegerTuple2.f1;
        }

        @Override
        public Tuple2<Integer, Integer> createAccumulator() {
            return Tuple2.of(0, 0);
        }

        /**
         * 累加方法，每来一行计算一行
         *
         * @param acc    累加器类型
         * @param score  分数
         * @param weight 权重
         */
        public void accumulate(Tuple2<Integer, Integer> acc, Integer score, Integer weight) {
            acc.f0 += score * weight; // 加权总和 = 分数1 * 权重1 + 分数2 *权重2 + ...
            acc.f1 += weight; // 权重和 = 权重1 + 权重2
        }
    }
}
